#include "lib/print.h"
#include "lib/str.h"
#include "mem/pmem.h"
#include "mem/vmem.h"
#include "proc/cpu.h"
#include "proc/initcode.h"
#include "memlayout.h"
#include "mem/mmap.h"
#include "riscv.h"
#include "fs/fs.h"
#include "fs/inode.h"
#include "fs/dir.h"
#include "fs/file.h"

/*----------------外部空间------------------*/
// in trampoline.S
extern char trampoline[];

// in swtch.S
extern void swtch(context_t* old, context_t* new);

// in trap_user.c
extern void trap_user_return();

// 内核页表
extern pgtbl_t kernel_pgtbl;
/*----------------本地变量------------------*/
// 进程数组
static proc_t procs[NPROC];

// 第一个进程的指针
static proc_t* proczero;

// 全局的pid和保护它的锁 
static int global_pid = 1;
static spinlock_t lk_pid;

static void proc_wakeup_one(proc_t* p);


// 申请一个pid(锁保护)
static int alloc_pid()
{
    int tmp = 0;
    spinlock_acquire(&lk_pid);
    assert(global_pid >= 0, "alloc_pid: overflow");
    tmp = global_pid++;
    spinlock_release(&lk_pid);
    return tmp;
}

// 释放锁 + 调用 trap_user_return
static void fork_return()
{
    static int first = 1;
    // 由于调度器中上了锁，所以这里需要解锁
    spinlock_release(&myproc()->lk);
    if (first)
    {
        first = 0;
        fs_init();

        proczero->cwd = path_to_inode("/");

        inode_lock(proczero->cwd);
        dir_print(proczero->cwd);
        inode_unlock(proczero->cwd);
        
        file_t* f=file_create_dev("./console", DEV_CONSOLE, 0);
        proc_t *p = myproc();
        p->filelist[0] = f;
        p->filelist[1] = file_dup(f);

        inode_lock(proczero->cwd);
        dir_print(proczero->cwd);
        inode_unlock(proczero->cwd);
    }
    trap_user_return();
}

// 返回一个未使用的进程空间
// 设置pid + 设置上下文中的ra和sp
// 申请tf和pgtbl使用的物理页
proc_t* proc_alloc()
{
    int pid=alloc_pid();
    proc_t *newproc=NULL;

    // 能否使用哈希优化？
    for(int i=0;i<NPROC;++i)
    {
        newproc=&procs[i];
        spinlock_acquire(&(newproc->lk));
        if(newproc->state==UNUSED)
        {
            break;
        }
        else
        {
            spinlock_release(&(newproc->lk));
            newproc=NULL;
        }
    }
    assert(newproc!=NULL, "proc_alloc: no proc available");

    // pid初始化
    newproc->pid=pid;

    // pagetable相关的初始化
    newproc->tf=pmem_alloc(false);
    newproc->ustack_pages=1;
    newproc->pgtbl=proc_pgtbl_init((uint64)newproc->tf);

    // 准备上下文切换
    memset(&(newproc->ctx),0,sizeof(newproc->ctx));
    newproc->ctx.ra=(uint64)fork_return;
    newproc->ctx.sp=newproc->kstack+PGSIZE;

    // 注意到正常结束时进程锁并未被释放，这样，我们在调用alloc创建一个新进程后不要忘了释放这个锁
    // 这是为了防止我们的进程在创建过程中被意外地调度
    return newproc;

}

// 释放一个进程空间
// 释放pgtbl的整个地址空间
// 释放mmap_region到仓库
// 设置其余各个字段为合适初始值
// tips: 调用者需持有p->lk
void proc_free(proc_t* p)
{
    // 释放之前为进程分配的trapframe
    if(p->tf)
        pmem_free((uint64)p->tf, false);
    p->tf = 0;

    // 可能需要清空页表
    if(p->pgtbl!=NULL){
        uvm_destroy_pgtbl(p->pgtbl);
    }

    p->pgtbl = NULL; // 避免悬空指针
    // 遍历并释放 mmap_region 链表中的所有节点
    mmap_region_t* current_region = p->mmap;
    while (current_region != NULL) {
        mmap_region_t* temp = current_region;
        current_region = current_region->next;
        mmap_region_free(temp); // 释放单个映射区域的内存
    }

    p->mmap = NULL; // 重置 mmap 指针
    // 将所有重要字段重置为初始状态，防止后续错误使用
    p->pid = -1;                // 设置为无效进程
    p->state = UNUSED;          // 将状态标记为未使用
    p->parent = NULL;           // 清空父进程引用
    p->exit_state = 0;          // 退出状态归零
    p->sleep_space = NULL;      // 清空睡眠等待空间
    p->heap_top = 0;            // 重置堆顶指针
    p->ustack_pages = 0;        // 清空用户栈页面数
    memset((void*)&(p->ctx), 0, sizeof(context_t)); // 清空内核上下文结构

}

// 进程模块初始化
void proc_init()
{
    struct proc *p;
    // 初始化全局 PID 自旋锁，保护 PID 分配
    spinlock_init(&lk_pid, "pid_lock");
    // 遍历每个进程，进行初始化
    for (p = procs; p < &procs[NPROC]; p++) {
        // 初始化每个进程的自旋锁
        spinlock_init(&p->lk, "process_lock");

        // 分配一页物理内存给内核栈
        char *pa = pmem_alloc(true);
        if (pa == 0)
            panic("pmem_alloc failed: insufficient memory for kernel stack");

        // 计算当前进程的内核栈虚拟地址
        uint64 va = KSTACK((int)(p - procs));
        
        // 将物理内存映射到内核页表
        vm_mappages(kernel_pgtbl, va, (uint64)pa, PGSIZE, PTE_R | PTE_W);
        p->kstack = va; // 更新进程的内核栈地址

        // 获取进程锁并释放资源
        spinlock_acquire(&p->lk);
        proc_free(p);  // 清理进程结构体，重置为 UNUSED 状态
        spinlock_release(&p->lk);
    }

    // 刷新TLB以应用新的页表映射
    kvm_inithart();
}

// 获得一个初始化过的用户页表
// 完成了trapframe 和 trampoline 的映射
pgtbl_t proc_pgtbl_init(uint64 trapframe)
{
    // 1. 分配页表
    pgtbl_t pgtbl = (pgtbl_t)pmem_alloc(true);
    if (pgtbl == 0) {
        panic("proc_pgtbl_init: unable to allocate page table\n");
    }
    // 2. 映射 trampoline
    vm_mappages(pgtbl, TRAMPOLINE, (uint64)trampoline, PGSIZE, PTE_R | PTE_X);
    // 3. 映射 trapframe
    vm_mappages(pgtbl, TRAPFRAME, trapframe, PGSIZE, PTE_R | PTE_W);
    return pgtbl;
}

/*
    第一个用户态进程的创建
    它的代码和数据位于initcode.h的initcode数组

    第一个进程的用户地址空间布局:
    trapoline   (1 page)
    trapframe   (1 page)
    ustack      (1 page)
    .......
                        <--heap_top
    code + data (1 page)
    empty space (1 page) 最低的4096字节 不分配物理页，同时不可访问
*/
void proc_make_first()
{
    uint64 page;
    proczero = proc_alloc();

    // ustack 映射 + 设置 ustack_pages 
    vm_mappages(proczero->pgtbl, TRAPFRAME - PGSIZE,(uint64)pmem_alloc(false), PGSIZE, PTE_W | PTE_R | PTE_U);
    proczero->ustack_pages = 1;

    // data + code 映射
    assert(initcode_len <= PGSIZE, "proc_make_first: initcode too big\n");
    page = (uint64)pmem_alloc(false);  // 分配一个物理页
    memmove((void*)page, initcode, initcode_len);  // 复制 initcode 数据
    vm_mappages(proczero->pgtbl, USER_BASE, page, PGSIZE, PTE_R | PTE_X | PTE_U | PTE_W);

    // 设置 heap_top
    proczero->heap_top = USER_BASE + PGSIZE;

    // 设置 mmap_region_t
    proczero->mmap = mmap_region_alloc(false);
    proczero->mmap->begin = MMAP_BEGIN;
    proczero->mmap->npages = (MMAP_END - MMAP_BEGIN) / PGSIZE;
    proczero->mmap->next = NULL;

    // tf字段设置
    proczero->tf->epc = USER_BASE;  // 用户程序的入口地址
    proczero->tf->sp = TRAPFRAME;   // 设置栈指针 sp 指向栈底
    // 设置proczero状态为RUNNABLE
    proczero->state = RUNNABLE;
    // 将proczero设置为这个cpu运行的进程
    cpu_t *c = mycpu();
    c->proc = proczero;
    spinlock_release(&proczero->lk);
}

// 进程复制
// UNUSED -> RUNNABLE
int proc_fork()
{

    // 分配一个新的进程空间，并获取其指针
    proc_t* new_proc = proc_alloc();
    assert(new_proc != NULL, "proc_fork: no proc available");

    // 设置子进程的 trapframe（用户态上下文），复制父进程的上下文
    // 子进程调用 fork 后，返回值存储在 a0 中，初始化为 0
    *(new_proc->tf) = *(myproc()->tf);  // 复制父进程的 trapframe
    new_proc->tf->a0 = 0;  // 子进程返回值是 0

    // 拷贝父进程的页表到子进程的页表
    // 复制父进程的堆区域、栈区域和内存映射区域
    uvm_copy_pgtbl(myproc()->pgtbl, new_proc->pgtbl, myproc()->heap_top, myproc()->ustack_pages, myproc()->mmap);

    // 创建新的 mmap 链表，复制父进程的 mmap 链表
    new_proc->mmap = mmap_region_alloc(false);  // 分配新的 mmap 区域
    mmap_region_t *new_region = new_proc->mmap;

    // 遍历父进程的 mmap 链表，逐个复制
    for (mmap_region_t *region = myproc()->mmap; region != NULL; region = region->next) {
        new_region->begin = region->begin;  // 复制映射区域的起始地址
        new_region->npages = region->npages;  // 复制映射区域的页面数

        // 如果父进程的当前映射区域还有下一个区域，继续分配新的节点
        if (region->next != NULL) {
            new_region->next = mmap_region_alloc(false);
        } else {
            new_region->next = NULL;  // 如果是最后一个节点，则设置为 NULL
        }
        new_region = new_region->next;  // 移动到下一个区域
    }

    // 设置新进程的堆顶和用户栈页数
    new_proc->heap_top = myproc()->heap_top;  // 复制堆顶
    new_proc->ustack_pages = myproc()->ustack_pages;  // 复制用户栈页数

    // 设置新进程的父进程为当前进程（myproc()）
    new_proc->parent = myproc();

        // 文件描述符的拷贝
    for(int i=0;i<FILE_PER_PROC;++i)
    {
        if(myproc()->filelist[i]!=NULL){
            new_proc->filelist[i]=file_dup(myproc()->filelist[i]);
        }
    }

    new_proc->cwd=inode_dup(myproc()->cwd);
    // 将新进程的状态设置为 RUNNABLE，表示可以被调度
    new_proc->state = RUNNABLE;

    // 释放新进程的锁
    spinlock_release(&(new_proc->lk));
    // 返回新进程的 PID，父进程通过这个 PID 判断是哪个子进程
    return new_proc->pid;

}

// 进程放弃CPU的控制权
// RUNNING -> RUNNABLE
void proc_yield()
{
    proc_t* current_proc = myproc(); // 获取当前运行的进程

    // 获取进程的锁，确保状态修改的原子性
    spinlock_acquire(&(current_proc->lk));

    // 将当前进程的状态设置为 RUNNABLE，表示它可以被重新调度
    current_proc->state = RUNNABLE;
    // 调用调度器，将控制权交回给 CPU 的调度进程
    proc_sched();

    // 释放锁，允许其他进程或调度器访问当前进程
    spinlock_release(&(current_proc->lk));

}

// 等待一个子进程进入 ZOMBIE 状态
// 将退出的子进程的exit_state放入用户给的地址 addr
// 成功返回子进程pid，失败返回-1
int proc_wait(uint64 addr)
{
    struct proc *np;
    int havechild;  // 记录当前进程是否拥有子进程
    int pid;
    struct proc *p = myproc();  // 获取当前进程（父进程）

    spinlock_acquire(&p->lk);  // 获取父进程的锁
    // 扫描全局进程组寻找已退出的子进程
    for(;;)
    {
        havechild = 0;

        // 遍历整个进程表
        for(np = procs; np < &procs[NPROC]; np++)
        {
            // 如果是子进程（父进程的子进程）
            if(np->parent == p)
            {
                // 确保子进程不是在退出中（防止在退出或切换中被调度）
                spinlock_acquire(&np->lk);  // 获取子进程的锁
                havechild = 1;

                // 找到退出的子进程
                if(np->state == ZOMBIE)
                {
                    pid = np->pid;
                    // 将退出的子进程的退出状态放入用户给定的地址 addr
                    uvm_copyout(p->pgtbl, addr, (uint64)&np->exit_state, sizeof(np->exit_state));

                    // 清理子进程资源
                    proc_free(np);

                    // 释放已经持有的锁
                    spinlock_release(&np->lk);
                    spinlock_release(&(myproc()->lk));
                    return pid;  // 返回子进程的 pid
                }

                // 释放子进程的锁
                spinlock_release(&np->lk);
            }
        }

        // 如果没有子进程，则返回 -1
        if(!havechild)
        {
            spinlock_release(&p->lk);  // 释放父进程的锁
            return -1;
        }
        // spinlock_release(&p->lk);  // 获取父进程的锁
        // 如果没有子进程进入 ZOMBIE 状态，父进程主动放弃 CPU，等待其他进程调度
        proc_sleep(p,&p->lk);  // 放弃当前进程的 CPU 控制权

        // 在重新调度后，继续遍历进程表
    }
}

// 父进程退出，子进程认proczero做父，因为它永不退出
static void proc_reparent(proc_t* parent)
{
    for (struct proc *child = procs; child < &procs[NPROC]; child++) {
        // 如果找到子进程，其父进程为 old_parent
        if (child->parent == parent) {
            child->parent = proczero;  // 将子进程的父进程改为 proczero
            proc_wakeup_one(proczero);  // 唤醒 proczero 检测子进程的退出
        }
    }
}

// 唤醒一个进程
static void proc_wakeup_one(proc_t* p)
{
    assert(spinlock_holding(&p->lk), "proc_wakeup_one: lock");
    if(p->state == SLEEPING && p->sleep_space == p) {
        p->state = RUNNABLE;
    }
}

// 进程退出
void proc_exit(int exit_state)
{
    struct proc *p = myproc();
    // 1号进程应该一直存在
    if(p == proczero)
        panic("proczero exiting");
    
    // 唤醒proczero进程
    spinlock_acquire(&proczero->lk);
    proc_wakeup_one(proczero);
    spinlock_release(&proczero->lk);

    // 如果原先的父进程退出，p->parent会变成proczero,原先父进程无法解锁，故先获取p->parent副本
    spinlock_acquire(&p->lk);
    proc_t *parent = p->parent;
    spinlock_release(&p->lk);

    // 唤醒父进程需要父进程的锁
    spinlock_acquire(&parent->lk);

    // 将当前进程的子进程交给proczero监管
    spinlock_acquire(&p->lk);
    proc_reparent(p);
    // 让当前进程的父进程回收这个退出进程
    proc_wakeup_one(p->parent);
    p->exit_state = exit_state;
    p->state = ZOMBIE;
    spinlock_release(&parent->lk);
    proc_sched();
    panic("zombie exit");
}

// 进程切换到调度器
// ps: 调用者保证持有当前进程的锁
void proc_sched()
{
    struct proc *current_proc = myproc();  // 获取当前运行的进程

    // 检查是否持有当前进程的锁
    if (!spinlock_holding(&current_proc->lk)) {
        panic("proc_sched: current process lock not held");
    }

    // 检查当前 CPU 是否释放了除当前进程锁之外的所有锁
    if (mycpu()->noff != 1) {
        panic("proc_sched: other locks not released");
    }

    // 确保当前进程不是 RUNNING 状态
    if (current_proc->state == RUNNING) {
        panic("proc_sched: current process still running");
    }

    // 确保中断已经关闭，避免在切换期间发生中断
    if (intr_get()) {
        panic("proc_sched: interrupts enabled");
    }

    // 保存当前 CPU 的中断开关状态
    int previous_interrupt_state = mycpu()->origin;

    // 进行上下文切换，保存当前进程的上下文，切换到调度器
    swtch(&current_proc->ctx, &mycpu()->ctx);

    // 恢复 CPU 的中断开关状态
    mycpu()->origin = previous_interrupt_state;

}

// 调度器
void proc_scheduler()
{
    struct proc *p;
    struct cpu *c = mycpu();
  
    // 设置当前CPU核心正在运行的进程为空,保证了内核进程不会被时钟中断打断
    c->proc = 0;
 
    for(;;)
    {
        intr_on();
        for(p = procs; p < &procs[NPROC]; p++) 
        {
            spinlock_acquire(&p->lk);
            if(p->state == RUNNABLE) 
            {
                //测试使用的printf
                //printf("proc %d is running\n", p->pid);
                p->state = RUNNING;
                c->proc = p;
                swtch(&c->ctx, &p->ctx);
                // c->proc=0表示运行在内核态
                c->proc = 0;
            }
            spinlock_release(&p->lk);
        }
    }

}

// 进程睡眠在sleep_space
void proc_sleep(void* sleep_space, spinlock_t* lk)
{
    proc_t* p = myproc();

    // 如果当前进程持有的锁与传入的锁不同，先获取当前进程的锁并释放传入的锁
    if(lk != &(p->lk))
    {
        spinlock_acquire(&p->lk);
        spinlock_release(lk);
    }

    p->sleep_space = sleep_space; // 设置睡眠原因标记
    p->state = SLEEPING;          // 设置进程为睡眠状态

    proc_sched(); // 放弃 CPU，进行调度，等待被唤醒

    // 被唤醒后清空睡眠原因标记
    p->sleep_space = 0;

    // 恢复原先的锁
    if(lk != &p->lk)
    {
        spinlock_release(&p->lk);
        spinlock_acquire(lk);
    }

}

// 唤醒所有在sleep_space沉睡的进程
void proc_wakeup(void* sleep_space)
{
    proc_t *p;
    for(p = procs; p < &procs[NPROC]; p++) 
    {
        spinlock_acquire(&p->lk);
        if(p->state == SLEEPING && p->sleep_space == sleep_space)
            p->state = RUNNABLE;
        spinlock_release(&p->lk);
    }
}